"""Run advection equation experiments from Section 3.2 of the paper."""
import sys
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.reproduct_paper.problems.advection import advection_1d_reference
from src.reproduct_paper.advection_solver import AdvectionSolver


def run_basic_test():
    """Run basic test with t_f=10, matching Figure 1 setup."""
    print("=" * 80)
    print("Running basic advection test (t_f=10, 10 time blocks)")
    print("=" * 80)

    problem = advection_1d_reference()

    # Parameters from paper: Figure 1
    # t_f = 10, Nb = 10 time blocks
    # Ne = 8 subdomains/block (Nx=4, Nt=2)
    # Q = 20x20 collocation points
    # M = 300 training parameters/subdomain
    # Rm = 1.0

    solver = AdvectionSolver(
        x_domain=problem["domain_x"],
        t_final=10.0,
        num_time_blocks=10,
        num_subdomains_x=4,
        num_subdomains_t=2,
        collocation_x=20,
        collocation_t=20,
        hidden_width=300,
        rand_magnitude=1.0,
        wave_speed=problem["wave_speed"],
        initial_condition=problem["initial_condition"],
        exact_solution=problem["exact_solution"],
        rng=np.random.default_rng(1),  # Fixed seed for reproducibility
    )

    results = solver.solve()
    overall_errors = solver.get_overall_errors()

    print(f"\nTotal training time: {results['total_time']:.2f} seconds")
    print(f"Average block time: {np.mean(results['block_times']):.2f} seconds")
    print(f"Maximum error: {overall_errors['max']:.6e}")
    print(f"RMS error: {overall_errors['rms']:.6e}")

    # Plot solution
    x_eval = np.linspace(0, 5, 100)
    t_eval = np.linspace(0, 10, 100)
    u_numerical = solver.evaluate_solution(x_eval, t_eval)

    # Compute exact solution
    u_exact = problem["exact_solution"](x_eval, t_eval)

    # Create plots
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))

    X, T = np.meshgrid(x_eval, t_eval, indexing='ij')

    # Plot numerical solution
    im1 = axes[0].contourf(X, T, u_numerical, levels=20, cmap='viridis')
    axes[0].set_xlabel('x')
    axes[0].set_ylabel('t')
    axes[0].set_title('locELM Solution')
    plt.colorbar(im1, ax=axes[0])

    # Plot error
    error = np.abs(u_numerical - u_exact)
    im2 = axes[1].contourf(X, T, error, levels=20, cmap='hot')
    axes[1].set_xlabel('x')
    axes[1].set_ylabel('t')
    axes[1].set_title('Absolute Error')
    plt.colorbar(im2, ax=axes[1])

    plt.tight_layout()
    output_path = Path(__file__).parent.parent / "experiments" / "advection_basic.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nSaved plot to: {output_path}")
    plt.close()

    return results, overall_errors


def run_subdomain_convergence():
    """Test convergence with varying number of subdomains (Figure 3a)."""
    print("\n" + "=" * 80)
    print("Testing convergence: varying number of subdomains (Figure 3a)")
    print("=" * 80)

    problem = advection_1d_reference()

    # Fixed: Q=20x20, M=300, Rm=1.0, t_f=10, Nb=10
    # Vary: Ne (number of subdomains/block)
    # Ne=2: (Nx,Nt)=(2,1), Ne=4: (2,2), Ne=8: (4,2)

    configs = [
        {"Ne": 2, "Nx": 2, "Nt": 1},
        {"Ne": 4, "Nx": 2, "Nt": 2},
        {"Ne": 8, "Nx": 4, "Nt": 2},
    ]

    results = []
    for config in configs:
        print(f"\nTesting Ne={config['Ne']}, Nx={config['Nx']}, Nt={config['Nt']}")

        solver = AdvectionSolver(
            x_domain=problem["domain_x"],
            t_final=10.0,
            num_time_blocks=10,
            num_subdomains_x=config["Nx"],
            num_subdomains_t=config["Nt"],
            collocation_x=20,
            collocation_t=20,
            hidden_width=300,
            rand_magnitude=1.0,
            wave_speed=problem["wave_speed"],
            initial_condition=problem["initial_condition"],
            exact_solution=problem["exact_solution"],
            rng=np.random.default_rng(1),
        )

        res = solver.solve()
        errors = solver.get_overall_errors()

        results.append({
            "Ne": config["Ne"],
            "max_error": errors["max"],
            "rms_error": errors["rms"],
            "time": res["total_time"],
        })

        print(f"  Max error: {errors['max']:.6e}, RMS error: {errors['rms']:.6e}")

    # Plot results
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))

    Ne_values = [r["Ne"] for r in results]
    max_errors = [r["max_error"] for r in results]
    rms_errors = [r["rms_error"] for r in results]

    ax.semilogy(Ne_values, max_errors, 'o-', label='Maximum error', linewidth=2, markersize=8)
    ax.semilogy(Ne_values, rms_errors, 's-', label='RMS error', linewidth=2, markersize=8)
    ax.set_xlabel('Number of sub-domains per time block (Ne)', fontsize=12)
    ax.set_ylabel('Error', fontsize=12)
    ax.set_title('Effect of number of sub-domains (Q=20x20, M=300)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    output_path = Path(__file__).parent.parent / "experiments" / "advection_subdomain_conv.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nSaved plot to: {output_path}")
    plt.close()

    return results


def run_collocation_convergence():
    """Test convergence with varying collocation points (Figure 3b)."""
    print("\n" + "=" * 80)
    print("Testing convergence: varying collocation points (Figure 3b)")
    print("=" * 80)

    problem = advection_1d_reference()

    # Fixed: Ne=8 (Nx=4, Nt=2), M=300, Rm=1.0, t_f=10, Nb=10
    # Vary: Qx=Qt (maintaining square grid)

    Q_values = [10, 15, 20, 25, 30]
    results = []

    for Q in Q_values:
        print(f"\nTesting Q={Q}x{Q}")

        solver = AdvectionSolver(
            x_domain=problem["domain_x"],
            t_final=10.0,
            num_time_blocks=10,
            num_subdomains_x=4,
            num_subdomains_t=2,
            collocation_x=Q,
            collocation_t=Q,
            hidden_width=300,
            rand_magnitude=1.0,
            wave_speed=problem["wave_speed"],
            initial_condition=problem["initial_condition"],
            exact_solution=problem["exact_solution"],
            rng=np.random.default_rng(1),
        )

        res = solver.solve()
        errors = solver.get_overall_errors()

        results.append({
            "Q": Q,
            "max_error": errors["max"],
            "rms_error": errors["rms"],
            "time": res["total_time"],
        })

        print(f"  Max error: {errors['max']:.6e}, RMS error: {errors['rms']:.6e}")

    # Plot results
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))

    Q_vals = [r["Q"] for r in results]
    max_errors = [r["max_error"] for r in results]
    rms_errors = [r["rms_error"] for r in results]

    ax.semilogy(Q_vals, max_errors, 'o-', label='Maximum error', linewidth=2, markersize=8)
    ax.semilogy(Q_vals, rms_errors, 's-', label='RMS error', linewidth=2, markersize=8)
    ax.set_xlabel('Collocation points per direction (Qx=Qt)', fontsize=12)
    ax.set_ylabel('Error', fontsize=12)
    ax.set_title('Effect of collocation points (Ne=8, M=300)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    output_path = Path(__file__).parent.parent / "experiments" / "advection_colloc_conv.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nSaved plot to: {output_path}")
    plt.close()

    return results


def run_training_param_convergence():
    """Test convergence with varying training parameters (Figure 3c)."""
    print("\n" + "=" * 80)
    print("Testing convergence: varying training parameters (Figure 3c)")
    print("=" * 80)

    problem = advection_1d_reference()

    # Fixed: Ne=8 (Nx=4, Nt=2), Q=20x20, Rm=1.0, t_f=10, Nb=10
    # Vary: M (training parameters/subdomain)

    M_values = [50, 100, 200, 300, 400, 500]
    results = []

    for M in M_values:
        print(f"\nTesting M={M}")

        solver = AdvectionSolver(
            x_domain=problem["domain_x"],
            t_final=10.0,
            num_time_blocks=10,
            num_subdomains_x=4,
            num_subdomains_t=2,
            collocation_x=20,
            collocation_t=20,
            hidden_width=M,
            rand_magnitude=1.0,
            wave_speed=problem["wave_speed"],
            initial_condition=problem["initial_condition"],
            exact_solution=problem["exact_solution"],
            rng=np.random.default_rng(1),
        )

        res = solver.solve()
        errors = solver.get_overall_errors()

        results.append({
            "M": M,
            "max_error": errors["max"],
            "rms_error": errors["rms"],
            "time": res["total_time"],
        })

        print(f"  Max error: {errors['max']:.6e}, RMS error: {errors['rms']:.6e}")

    # Plot results
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))

    M_vals = [r["M"] for r in results]
    max_errors = [r["max_error"] for r in results]
    rms_errors = [r["rms_error"] for r in results]

    ax.semilogy(M_vals, max_errors, 'o-', label='Maximum error', linewidth=2, markersize=8)
    ax.semilogy(M_vals, rms_errors, 's-', label='RMS error', linewidth=2, markersize=8)
    ax.set_xlabel('Training parameters per sub-domain (M)', fontsize=12)
    ax.set_ylabel('Error', fontsize=12)
    ax.set_title('Effect of training parameters (Ne=8, Q=20x20)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    output_path = Path(__file__).parent.parent / "experiments" / "advection_trapar_conv.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nSaved plot to: {output_path}")
    plt.close()

    return results


def run_random_magnitude_test():
    """Test effect of random coefficient magnitude (Figure 4)."""
    print("\n" + "=" * 80)
    print("Testing effect of random coefficient magnitude (Figure 4)")
    print("=" * 80)

    problem = advection_1d_reference()

    # Fixed: Ne=8 (Nx=4, Nt=2), Q=20x20, M=300, t_f=10, Nb=10
    # Vary: Rm

    Rm_values = [0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0]
    results = []

    for Rm in Rm_values:
        print(f"\nTesting Rm={Rm}")

        solver = AdvectionSolver(
            x_domain=problem["domain_x"],
            t_final=10.0,
            num_time_blocks=10,
            num_subdomains_x=4,
            num_subdomains_t=2,
            collocation_x=20,
            collocation_t=20,
            hidden_width=300,
            rand_magnitude=Rm,
            wave_speed=problem["wave_speed"],
            initial_condition=problem["initial_condition"],
            exact_solution=problem["exact_solution"],
            rng=np.random.default_rng(1),
        )

        res = solver.solve()
        errors = solver.get_overall_errors()

        results.append({
            "Rm": Rm,
            "max_error": errors["max"],
            "rms_error": errors["rms"],
            "time": res["total_time"],
        })

        print(f"  Max error: {errors['max']:.6e}, RMS error: {errors['rms']:.6e}")

    # Plot results
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))

    Rm_vals = [r["Rm"] for r in results]
    max_errors = [r["max_error"] for r in results]
    rms_errors = [r["rms_error"] for r in results]

    ax.semilogy(Rm_vals, max_errors, 'o-', label='Maximum error', linewidth=2, markersize=8)
    ax.semilogy(Rm_vals, rms_errors, 's-', label='RMS error', linewidth=2, markersize=8)
    ax.set_xlabel('Random coefficient magnitude (Rm)', fontsize=12)
    ax.set_ylabel('Error', fontsize=12)
    ax.set_title('Effect of random coefficients (Ne=8, Q=20x20, M=300)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    output_path = Path(__file__).parent.parent / "experiments" / "advection_randmag.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nSaved plot to: {output_path}")
    plt.close()

    return results


if __name__ == "__main__":
    # Run all experiments
    basic_results, basic_errors = run_basic_test()
    subdomain_results = run_subdomain_convergence()
    colloc_results = run_collocation_convergence()
    trapar_results = run_training_param_convergence()
    randmag_results = run_random_magnitude_test()

    print("\n" + "=" * 80)
    print("All experiments completed!")
    print("=" * 80)
